# coding=utf-8
# 获取指数涨跌
import os
import re
import sys
import traceback

import akshare as ak
import pandas
import pandas as pd
from loguru import logger

from models.stock_model import StockNumber, DayInfo
from mylib import update_all_csv
from mylib.download_all import analysis_stock
from mylib.mycsv import sort_csv


def run(d=1):
    if d:
        sz_index_df = ak.index_zh_a_hist(symbol="000001", period="daily")
        df_sorted = sz_index_df.sort_values(by='日期', ascending=False)
        df_sorted.to_csv("shanghai_index.csv", index=False)
    else:
        df_sorted = pd.read_csv('shanghai_index.csv')
    today_date = eval(str(list(df_sorted['日期'])[0]).replace('-', ''))
    return today_date


def get_vol2(N, today_date, sn):
    sc = f'stocks/{sn.ts_code}.csv'
    if not os.path.exists(sc):
        return None, None
    df = pandas.read_csv(sc)
    if str(df['trade_date'][0]) != str(today_date):
        analysis_stock(sn)
        df = pandas.read_csv(sc)
    result_arr = list()
    d_pre = None
    d_arr = []
    for row in df.index:
        if len(result_arr) > N:
            break
        d = DayInfo(sn, df.loc[row])
        if d_pre and d_pre.pct_chg > 3 and d_pre.vol > d.vol * 2:
            cur_k = d_pre.trade_date
            result_arr.append((cur_k, len(d_arr) - 1))
            d_arr = []
        d_arr.append(d)
        d_pre = d
    res = re.findall('stocks/(.*).csv', sc)
    link_code_arr = res[0].split('.')
    link_code = f'{link_code_arr[1]}{link_code_arr[0]}'
    hyperlink = f'"https://xueqiu.com/S/{link_code}"'
    date_arr_hyp = f'{res[0]},=HYPERLINK({hyperlink})'
    return date_arr_hyp, result_arr


def get_all_stock_csv_path():
    for root, dirs, files in os.walk('stocks'):
        return [os.path.join(root, item) for item in files]


if __name__ == '__main__':
    try:
        if len(sys.argv) == 2:
            N = int(sys.argv[1])
        else:
            N = 3
    except Exception as e:
        logger.error(e)
        N = 3
    # update_all_csv.run()
    while N > 1:
        D = True
        today_date = run(D)
        df = pd.read_csv('cal_ops/all.csv')
        full_path_txt = f'{today_date}_get_vol2_N{N}.txt'
        full_path_csv = f'{today_date}_get_vol2_N{N}.csv'
        if os.path.exists(full_path_txt):
            os.remove(full_path_txt)
        logger.add(full_path_txt, format='{message}')
        n_arr = []
        ts_code_arr = []
        result_dict = {}
        cnt = 0
        with open(full_path_csv, 'w', encoding='utf-8') as aa:
            aa.write(f'AA,BB,CC,DD,EE,FF,GG,HH')
            for row in df.index:
                cnt += 1
                # if cnt > 10:
                #     break
                sn = StockNumber(df.loc[row])
                if 'ST' in sn.name:
                    continue
                try:
                    date_arr_hyp, result_arr = get_vol2(N, today_date, sn)
                    sum_days = 0
                    days_str = ''
                    days_str2 = f'({N})'
                    cur_len = result_arr[0][1]
                    for indx, item in enumerate(result_arr):
                        tdate, days = item
                        days_str += f'({tdate}:{days})-'
                        if indx:
                            days_str2 += f'-{days}'
                            sum_days += days
                    avg_d = round(sum_days / N, 2)
                    if avg_d > cur_len:
                        w_msg = f'\n{cur_len},{avg_d},{days_str2},{sn.industry},{sn.name},{date_arr_hyp},{days_str}'
                        aa.write(w_msg)
                except Exception as e:
                    print(e, traceback.format_exc())
        sort_csv(full_path_csv, ['AA', 'BB'], [True, True])
        if os.path.exists(full_path_txt):
            os.remove(full_path_txt)
        N -= 1
